perf: stop holding the fair pool lock across blocking memory calls by dwsmith1983 · Pull Request #5613 · apache/datafusion-comet

dwsmith1983 · 2026-09-01T15:16:58Z

Which issue does this PR close?

No dedicated issue. Adjacent to #5494, which made task-shared pools the norm and so widened the blast radius of this lock.

Rationale for this change

CometFairMemoryPool held its internal mutex across the JNI acquire and release calls into Spark's TaskMemoryManager. That call can block for a long time while Spark spills other consumers, and while it blocked, every other native thread sharing the task pool sat behind the lock, including plain releases that needed nothing from the JVM. The sibling unified pool already avoids this.

What changes are included in this PR?

The admission check and the reservation now happen as one short locked step (the fair limit couples used bytes and the consumer count, so this part genuinely needs mutual exclusion), then the blocking JVM call runs with no lock held, and the reservation rolls back if the JVM declines, grants partially, or the call panics. Fairness semantics are unchanged: concurrent grows still cannot jointly exceed pool_size divided by the consumer count, and registering a new consumer still only blocks further growth rather than clawing back existing reservations. Going fully lock-free like the unified pool was considered and rejected, since separate atomics would let a register or unregister slip between reading the count and committing the reservation.

Spark's ExecutionMemoryPool removes a task's accounting entry the moment its balance reaches zero, and an acquire parked inside Spark indexes that entry when it wakes, so a release that zeroes the balance while another acquire is parked crashes the waiter. With the lock no longer held across the call that window is real, so the pool holds a one byte anchor for as long as it lives: no request ever carries it and the balance never reaches zero mid task. Creating the pool makes no JVM call. The anchor is taken by the first grow that passes the pool's own fair-limit check and before that grow's real request, so a plan that never allocates natively never touches Spark's memory manager and never counts as an active task there, and a grow the pool rejects locally never reaches Spark while the anchor is missing. Spark declines the byte outright, with a zero grant, for a task already at its share, which a task holding JVM shuffle pages when its second native plan grows can reach; the pool then runs without its anchor and each grow retries the byte as a request of its own until it is held, so the task spills on the short grant instead of failing. That costs one extra JNI call per grow only while the anchor is missing. A grow whose anchor request fails or panics rolls its reservation back the same way a failed request does. The task-shared pool registry creates the pool outside its process-wide lock and re-checks on insert; a concurrent loser drops its own pool, which held nothing.

Two side effects worth naming. The JNI boundary moved behind a small internal trait so the pool can be tested without a live JVM (neither pool had any tests before). And the old code could deadlock if a blocked acquire ever re-entered the pool on the same thread via a spill callback, since the lock was held across the call; that hazard is gone by construction.

How are these changes tested?

Twenty-seven pool tests and nine registry tests, all new. The stub models Spark 4.1.3's ExecutionMemoryPool: the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, and notifyAll on release, with a bounded wait that fails a test instead of hanging. Covered: fairness rejection without reaching Spark, limit tightening on register, partial-grant rollback with the excess returned, acquire failure and panic rollback, zero-size no-op, over-shrink panic, a parked acquire not stalling a concurrent release, an acquire below its minimum share being granted after the holder frees memory in full, a release already on its way when a late acquire parks, a sibling consumer of the same task freeing its last page while an acquire is parked, a second acquire not queuing behind a parked first one inside the pool (Spark's TaskMemoryManager serializes a task's acquires on its own monitor, so this pins the Rust side only), an exact-fit first grant being accepted, the anchor taken on the first grow and returned once at drop, a pool that never grows never calling Spark, an over-limit grow without an anchor never reaching Spark, the first grow failing cleanly when the bridge errors on the anchor, a declined anchor leaving a usable pool that reports short grants and takes the byte on the first grow after the share frees with no second request once held, two grows retrying the missing anchor at once keeping exactly one byte, the registry creating outside its lock with neither create touching Spark, and an eight-thread stress test asserting the accounting never exceeds the fair limit and nets to zero. All looped 30 times in debug and release with no failures. Full core crate suite passes, clippy with warnings denied and fmt are clean.

Microbenchmarks against a real Spark 4.1.3 TaskMemoryManager over an off-heap UnifiedMemoryManager, driven through the JNI bridge with base and head in one binary (details and the full table in the discussion): at 4 and 8 native threads the release p99 falls from tens of microseconds to single digits because a release no longer waits behind an in-flight acquire, mean grow latency and throughput improve 1.5x to 2.5x under contention, grow p99 is unchanged since that is Spark's own work, and errors and final balances are zero on every row.

CometFairMemoryPool held its mutex across the JNI calls into Spark task memory manager, which can block for seconds while Spark spills, so every native thread sharing a task pool serialized behind whichever thread was acquiring. The fairness check and the reservation are now one short locked step, the blocking call runs unlocked, and the reservation rolls back if the JVM fails to back it or the call panics. Fairness semantics are unchanged: concurrent grows still cannot jointly exceed pool_size divided by the consumer count. The JNI boundary moved behind a small trait so the pool finally has tests: fairness rejection, limit tightening on register, partial-grant rollback, panic rollback, a blocking test that took ten seconds on the old code and 30ms now, and an eight-thread stress test.

sunchao

Reviewed 626bc395091f313ae93b8867c1cb1c846dc7fc53 against 8729f6e6adf7091e18a48670e790d4ba8fd41e51. I found one P2 in the new concurrent release path, detailed inline.

The focused Spark memory-pool component probe reproduced the missing-task exception. No full Comet native/JNI query suite ran. CI, CodeQL, and Delta Contrib Build Gate currently report action_required, with no test checks recorded.

For this performance change, please include matched BASE/HEAD microbenchmarks with one and multiple native threads, full and partial grants, and consumer registration changes. Report completed operations, grow/release latency, error counts, and final Rust/Spark balances. The parked stub test establishes lock behavior but does not measure production JNI throughput.

sunchao · 2026-09-02T02:43:29Z

+                state.used -= subtractive;
            }


[P2] Keep waiting Spark acquisitions registered during a full release

Could you handle Spark's same-task wait/release contract before allowing this release to overlap an acquire? With fair_unified, two native consumers can now enter the same CometTaskMemoryManager concurrently. In a 100-unit executor pool, let another task hold 90 and this task hold 10. This task's next 10-unit grow waits below Spark's 1/(2N) minimum. Freeing its last 10 units on another native thread removes its entry from ExecutionMemoryPool.memoryForTask and wakes the grower. The grower then indexes the removed entry and throws NoSuchElementException: key not found. Spark's release bypasses the task monitor held by the waiting acquire, so that monitor does not prevent this interleaving. The Rust provisional reservation does not keep Spark's entry alive.

I reproduced the failure using unchanged Spark 3.5.9 pool source with only logging/annotation/memory-mode scaffolding. A scheduling control modeling the previous serialization completed after the other task freed its memory. The relevant map lifecycle is also present in 4.0.4 source. This was a component probe plus JNI source tracing, not a full Comet query reproduction. Please make full releases safe while grows are pending and add a regression that exercises Spark's memory manager.

Confirmed against the pool source, thanks for the repro. Blocking the release until the in-flight acquires drain turned out to deadlock in testing, since the parked acquire can be waiting on exactly the memory that release frees. So the fix defers instead: a release that would zero the task's balance while acquires are in flight frees n-1 bytes right away (that is what wakes the waiter) and holds the last byte, which the final completing acquire pays off. At most one byte is ever deferred and it always settles once the acquires finish. The test stub now models the entry lifecycle (created on acquire, removed at zero, a woken waiter fails if the entry is gone) and reproduced this crash before the fix. It also turned out one of our existing tests was exercising the same broken pattern.

The new deferral fixes the already-pending case, but I can still reproduce the same missing-task failure through a late-arriving acquire on current head eb410e51.

At current lines 252-259, plan_release can see pending_acquires == 0, schedule the whole balance, and drop the state lock before release reaches Spark. A new try_grow can then increment pending_acquires and park in Spark while the old balance is still present. The already-planned release removes the task entry, and the waiter resumes with NoSuchElementException.

I reproduced this with the exact head's production state machine and a gated bridge: hold 10, plan and pause the full release, start and park a grow of 10, then resume the release. The waiter panics with key not found: task entry removed while acquire waited. paying_deferred fences only deferred payments, so could you also coordinate ordinary zeroing releases with newly starting acquires?

Also covered by db1f1bc. With the anchor there is no zeroing release left to coordinate: once any reservation exists the anchor is held, so a full release leaves the task at one byte and the entry survives.

Two related windows came out of review and are closed in the same commit. A grant that covers the request but not the extra byte is handed back as a short grant rather than running unanchored, and every acquire that starts while the anchor request is still parked carries its own extra byte, with the first full grant keeping it and later ones returning theirs. Your gated repro is pinned as late_acquire_survives_a_release_already_on_its_way, alongside a test for the concurrent in-flight case; both failed on eb410e5 with key not found and pass now. 30 loops each in debug and release are clean.

The earlier build-gate failure compared dylib sizes on a change confined to fair_pool.rs, so it looks like the size check rather than this branch.

One more window closed in 353541d, found while probing the anchor bootstrap with a third task. Before the anchor lands, a short grant is handed back whole; if another task's entry disappears in that gap (raising this task's minimum share) and a sibling acquire of this task then parks, the rollback release zeroes the entry under it. A small bootstrap mutex now serializes anchor carriers from the bridge acquire through the rollback release. Releases never take it, so a parked carrier cannot starve anyone, and non-carriers only exist once the anchor is held. Pinned as short_grant_rollback_cannot_land_under_a_sibling_parked_in_spark, which failed with key not found before the change.

A release that would zero the JVM-side balance while other acquires are still in flight frees all but one byte immediately and holds the last byte until the in-flight acquires complete. Spark drops the task's accounting entry when its balance hits zero, so a parked acquire waking after that point indexes a missing entry and fails. Blocking the release instead can deadlock because the parked acquire may be waiting for the very memory the release frees. The stub task memory now models the entry lifecycle so the regression is covered.

dwsmith1983 · 2026-09-03T10:56:51Z

@sunchao any more feedback here?

sunchao

Reviewed eb410e5175417256f0223deaee1b122d760ae077 against ef62b46306e925bc51e7d7f29922c1870eb729e7. I found two actionable issues in the revised release/acquire protocol: the held-back byte can leave a valid Spark acquire waiting indefinitely, and an ordinary full release can still be overtaken by a newly starting acquire. The first is inline below. I added the second as a follow-up on the existing unresolved thread because it is the same missing-task failure under a different ordering.

Local validation: the focused native fair-pool suite passed 12/12 tests with 241 filtered out. An exact-head component probe reproduced the ordinary-release ordering failure. A probe using unmodified Spark 4.1.3 ExecutionMemoryPool reproduced both the missing-key failure and the one-byte wait. These are component probes, not a full Comet native/JNI query run.

Current CI has 63 successful checks, 9 skipped checks, and one failed Delta build gate. Its log fails the contrib-enabled-versus-default libcomet size invariant; this PR changes only fair_pool.rs.

sunchao · 2026-09-03T20:24:46Z

+            .jvm_held
+            .checked_sub(bytes)
+            .expect("released more bytes than the JVM side holds");
+        if bytes > 0 && state.jvm_held == 0 && state.pending_acquires > 0 {


[P1] Avoid retaining the byte a waiting acquire may need

Could you avoid withholding one byte until finish_acquire? This can deadlock the default off-heap pool. In a 1 GiB Spark execution pool, let another task hold 900 MiB, let this task hold 100 MiB, and let a second native consumer for this task request 100 MiB. Spark parks that request because this task is below its 1/(2N) minimum share. When the holder frees 100 MiB, this branch sends only 100 MiB - 1. On wake, Spark computes toGrant = 100 MiB - 1; because that is short of the request and curMem + toGrant = 100 MiB is still below 256 MiB, it waits again. The deferred byte is paid only by finish_acquire, which cannot run while this acquire is waiting.

I reproduced this against unmodified Spark 4.1.3 ExecutionMemoryPool: retaining one byte left the grower in WAITING, and freeing one additional byte from the other task let it complete. The new stub test misses this because it grants the full request after any release without reapplying Spark's free-memory and minimum-share checks. Could you preserve the task entry without withholding capacity needed by the waiter?

Fixed in db1f1bc. The deferral is gone: every release now goes to the JVM whole, so a parked acquire sees the full freed amount and Spark's minimum-share check passes. The task entry is kept alive by a permanent anchor instead. The first acquire asks for one extra byte and the pool holds it until it drops, so the balance never reaches zero mid task. That means the task stays in Spark's active-task set for the pool's lifetime and retains one byte, which the header comment now states.

Your scenario is pinned as min_share_wait_is_granted_after_the_holder_frees_its_memory_in_full. The stub now models ExecutionMemoryPool.acquireMemory from 4.1.3 with the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, and notifyAll on release, with a bounded wait that fails the test instead of hanging. It failed on the previous head with the waiter timing out and passes now.

Spark's ExecutionMemoryPool removes a task's entry when its balance hits zero, and an acquire parked inside Spark indexes that entry on wake. The previous fix held one byte back from a zeroing release while acquires were in flight and paid it back later. That withheld byte starved Spark's minimum share check, which grants a parked request only when the freed bytes cover it in full, so the waiter slept forever; and a release planned before a late acquire parked could still remove the entry. The pool now asks for one extra byte with every acquire that starts before the anchor lands, keeps the first one until the pool drops, and sends every release whole. The task's balance never returns to zero mid task, so both failure modes are impossible by construction. The test double now models ExecutionMemoryPool's entry lifecycle, share rules and wait loop, and covers the minimum share scenario, the late acquire race and the in flight anchor race.

…e task entry Until the anchor lands, a short grant is handed back to Spark whole. In the gap between that grant returning and its rollback release landing, a sibling acquire of the same task can enter Spark and park if another task has left the pool meanwhile, and the release then removes the task's entry under it. Carriers now run one at a time from the bridge acquire through the rollback release. Releases never take that lock, so a parked carrier cannot hold up the release it waits for, and Spark serializes a task's acquires anyway.

dwsmith1983 · 2026-09-04T04:23:02Z

@sunchao both of your findings are addressed on the branch head (353541d), with replies on each thread. Ready for another look whenever you have time.

dwsmith1983 · 2026-09-09T00:23:10Z

@sunchao the anchor redesign replacing the deferral is pushed, with the Spark-faithful stub and your P1 and late-arrival cases as tests. Ready for another look.

andygrove

The anchor invariant this design rests on does not hold during the window in which the anchor is being established, and bootstrap only covers half of that window. In native/core/src/execution/memory_pools/fair_pool.rs around line 345 every carrier takes bootstrap and holds it across the bridge acquire, so until the anchor lands the pool is back to one blocking JNI call at a time, which is the serialization this PR set out to remove. That lock also serializes only this pool's own carriers. CometTaskMemoryManager.internal is the task-wide TaskMemoryManager, and CometUnifiedShuffleMemoryAllocator is another off-heap consumer on the same task, so if it frees its last page while a pre-anchor carrier is parked inside Spark, memoryForTask loses the entry and the waiter hits the same NoSuchElementException the anchor exists to prevent. Could the anchor instead be taken once as a standalone one byte acquire when the pool is set up, so that no ordinary request ever carries it and Anchor::Requested, bootstrap, and the surplus path all go away?

The extra byte on the request path also turns grants Spark would have satisfied into failures. exact_fit_grant_without_the_anchor_is_rejected_as_short pins the case where Spark backs additional but not additional + 1 and the whole grant is handed back as short. grow at line 272 is try_grow(...).unwrap(), so a caller on that path gets a panic rather than a spill. It is a knife edge, but it is a knife edge that only shows up under memory pressure. Is that acceptable, or is it another argument for keeping the anchor off the request path entirely?

#5466 rewrites the same admission check in this file and moves the fair-share basis from the pool-wide state.used to reservation.size(). DataFusion does not update reservation.size() until try_grow returns, so that change would undercut the reasoning in the comment at line 307 that reserving under the lock is what makes the unlocked window safe. Whichever of the two lands second needs a deliberate rebase rather than a textual merge.

Worth noting that CI has not run on b7007ea9 at all, since the workflows are still sitting at action_required. The microbenchmarks sunchao asked for in the first round are also still open, so the only evidence behind the performance claim right now is the stub timing.

…t path CometFairMemoryPool::try_new acquires the one byte anchor from Spark before the pool is usable and fails construction if Spark declines it, so every live pool holds the anchor for its whole life and no ordinary request carries it. This removes Anchor::Requested, the bootstrap mutex that serialized pre-anchor acquires, the surplus release, and the exact-fit rejection; try_grow asks for exactly the requested bytes again. A sibling consumer of the same task freeing its last page can no longer drop the task's memoryForTask entry under a parked acquire, since the anchor keeps the balance above zero. Because construction can now park inside Spark, acquire_task_shared_pool runs the create closure without the process-wide registry lock and keeps the first pool registered when two plans of one task create at once. create_memory_pool returns CometResult and createPlan propagates the error.

dwsmith1983 · 2026-09-09T14:43:24Z

@andygrove the anchor is now a standalone one byte acquire at pool setup (d004f04). try_new takes it before the pool exists and fails construction if Spark grants less than a byte or errors, so every live pool holds the anchor from birth to drop and no request carries it: Requested, the bootstrap mutex, the surplus release, and the exact-fit rejection are gone, and try_grow asks for exactly additional again with the reviewed reserve-under-lock, call-unlocked, roll-back shape. Your sibling-consumer case is a test now: the stub models another consumer of the same task freeing its last page while an acquire is parked, and the entry survives because the anchor byte is held. So is the serialization point: a second acquire no longer queues behind a parked first one.

One consequence worth calling out: construction can now park inside Spark, and acquire_task_shared_pool used to run the create step under the process-wide registry lock, which would have stalled plan creation for every task on the executor. It now creates outside the lock and re-checks on insert; a concurrent loser drops its own pool and its anchor goes back over JNI after the lock is released, with a test through the real registry.

On grow, the unwrap predates this branch and is DataFusion's contract for the infallible path; with the anchor off the request path the only panic left there is a short grant a plain request would also have hit.

On #5466: it moves the admission basis from state.used to reservation.size(), which DataFusion updates only after try_grow returns, so two concurrent grows of one reservation would both pass the check before either lands. Whichever lands second needs a deliberate rebase: keep the reserve-under-lock step and apply the per-reservation limit to it, and rewrite the two comments in try_grow and shrink that explain the state.used basis. Nothing in this commit changes that shape.

dwsmith1983 · 2026-09-09T14:45:05Z

@sunchao the microbenchmarks from your first pass, against a real Spark task memory manager rather than the stub: the harness launches a JVM in-process through the jni crate's invocation feature, builds a Spark 4.1.3 TaskMemoryManager over an off-heap UnifiedMemoryManager with a 1 GiB pool, wraps it in CometTaskMemoryManager, and drives base (main's fair_pool.rs compiled as a sibling module) and head through the real JNI bridge in one binary. Three recorded rounds after warmup, 20,000 operations per thread. Latencies are mean / p50 / p99 in microseconds. Fair-limit rejects and other errors were zero on every row; Rust used and the Spark balance netted to zero after quiesce on every row, with head holding its one anchor byte until drop and zero after.

scenario	threads	ops/s base to head	grow base to head	release base to head	short grants base / head
full grants	1	1.52M to 1.92M (1.26x)	0.7/0.7/0.9 to 0.5/0.5/0.7	0.4/0.4/0.5 to 0.4/0.3/0.5	0 / 0
partial grants	1	1.15M to 1.57M (1.36x)	0.8/0.7/7.0 to 0.6/0.5/3.0	0.4/0.4/1.2 to 0.4/0.4/0.5	1941 / 1941
registration churn	1	1.89M to 1.94M (1.03x)	0.5/0.5/1.5 to 0.5/0.5/1.2	0.4/0.4/1.1 to 0.4/0.4/0.6	0 / 0
full grants	4	0.92M to 1.37M (1.50x)	4.4/1.3/32.4 to 3.6/1.2/26.5	4.2/1.0/31.8 to 0.8/0.6/3.7	0 / 0
partial grants	4	0.72M to 1.34M (1.86x)	5.2/1.5/43.9 to 3.1/1.1/40.3	7.5/1.4/67.4 to 0.8/0.7/5.8	29314 / 23544
registration churn	4	0.97M to 1.60M (1.65x)	4.1/0.7/37.3 to 3.0/1.1/26.2	4.0/0.5/37.3 to 0.7/0.6/1.5	0 / 0
full grants	8	0.61M to 1.40M (2.27x)	12.9/3.7/83.4 to 7.2/1.2/93.1	12.9/4.0/83.3 to 0.8/0.6/2.5	0 / 0
partial grants	8	0.49M to 1.23M (2.49x)	16.5/12.4/93.5 to 7.1/1.2/129.5	13.4/1.8/89.2 to 0.7/0.7/1.8	110853 / 78614
registration churn	8	0.62M to 1.40M (2.27x)	12.9/2.6/92.1 to 7.2/1.1/91.9	12.8/2.5/92.0 to 0.8/0.6/3.0	0 / 0

The release path is where the change shows: at 4 and 8 threads its p99 drops from tens of microseconds to single digits because a release no longer waits behind an in-flight acquire. Grow p99 is unchanged, since that is Spark's own work under the call, while mean grow latency and throughput improve 1.5x to 2.5x under contention. Head sees fewer short grants in the partial scenarios because it cycles past the release faster. Registration churn counts only depend on run length. The harness needs a JDK and a full Spark classpath so it is not a criterion bench in tree; the source and raw output are available if wanted. macOS arm64, single machine.

andygrove

This is the shape I was hoping for, and it is a lot simpler than what it replaces. Anchor::Requested, the bootstrap mutex, the surplus release and the exact-fit rejection are all gone, and try_grow is back to reserve-under-lock, call-unlocked, roll-back-on-shortfall with no extra byte riding along. Moving the create step out from under the process-wide registry lock was the right call and I would not have thought to check it until it bit someone: with the anchor acquire able to park, holding TASK_SHARED_MEMORY_POOLS across it would have stalled plan creation for every task on the executor. The std::ptr::eq guard in TaskSharedMemoryPool::drop already covers the loser, so a concurrent create hands its anchor back without deregistering the winner. Thanks also for running the benchmarks against a real TaskMemoryManager rather than the stub; the release p99 collapsing from tens of microseconds to single digits at 4 and 8 threads is the number that actually demonstrates the claim, and it is the one the stub could not show.

One thing came out of the move that I do not think we should ship as is.

A declined anchor now fails the task instead of spilling

with_bridge returns CometError::Internal when Spark grants less than the byte, and createPlan propagates it, so the task dies before any work is attempted. Walking Spark's ExecutionMemoryPool.acquireMemory with numBytes = 1: maxToGrant = min(1, max(0, maxMemoryPerTask - curMem)), so a task whose curMem has already reached maxMemoryPerTask gets maxToGrant = 0, and the park guard curMem + toGrant < minMemoryPerTask is false for it because minMemoryPerTask is half of maxMemoryPerTask. It does not wait. It returns 0, and the pool refuses to exist.

That is reachable without anything exotic. The task only needs to be holding its share of execution memory at the moment the first native plan is created, and CometTaskMemoryManager.internal is the task-wide manager, so CometUnifiedShuffleMemoryAllocator and any Spark-side operator in a mixed plan count toward curMem. The concrete sequence I would worry about is a task that runs one native plan, closes it so the shared pool drops, allocates JVM shuffle pages, and then creates a second native plan: lookup misses, try_new runs again, and Spark declines. On main that situation produces a short grant on the first real try_grow and DataFusion spills. On this branch it is a task failure with no spill available as an escape, and the message reads like an internal error rather than memory pressure.

Could a declined anchor be non-fatal instead? Keep a flag for "anchor not held" and re-attempt a standalone one byte acquire at the top of try_grow while it is unset. That keeps everything this commit won: the anchor never rides on a real request, so there is no exact-fit hazard, no surplus to hand back and no bootstrap mutex, and the request path stays plain. The cost is one extra JNI call per try_grow only while the anchor is missing, which is the rare path by construction. The window where the pool holds reservation bytes without the anchor is then real but bounded, and it is strictly better than the window on main, which is the whole first acquire. A test for the declined-at-setup case would pin it either way; right now the stub only covers the grant.

The parking case reads fine to me and I convinced myself it cannot deadlock: a task parked in try_new holds nothing, so a pool fully consumed by tasks that are all parked in construction is not a reachable state, and maybeGrowPool covers the storage-borrow side.

Two smaller things. with_bridge's short-grant branch returns without releasing granted, which is dead today because ANCHOR_BYTES is 1 so a short grant is necessarily 0, but it would leak silently if that constant ever moved; a debug_assert or a one-line comment tying the two together would be cheap. And bridge.acquire at construction is the one bridge call not wrapped in catch_unwind. That is correct as written, since there is no optimistic reservation to roll back and try_unwrap_or_throw catches the panic at the JNI boundary, but the comment above the catch_unwind in try_grow explains why it is needed there without saying why it is not needed here.

Your read on #5466 matches mine, including that the per-reservation basis has to be applied to the reserve-under-lock step rather than replacing it, and that the two comments explaining the state.used basis need rewriting rather than merging. Worth agreeing an order with @peterxcli so whoever is second plans for it.

Spark grants zero bytes to a task already at its share, which a task holding JVM shuffle pages when its second native plan is created can reach. Failing pool construction there fails the task with no spill available, where main produces a short grant and spills. Create the pool without the anchor when Spark declines it, and have each grow retry a standalone one byte acquire first while it is missing, so the anchor never rides on a real request and the task spills on the short grant instead. A second grow racing the retry hands its byte back; drop releases the anchor only when held. A compile-time assert ties the zero grant reasoning to the one byte anchor size.

dwsmith1983 · 2026-09-10T12:05:47Z

Could a declined anchor be non-fatal instead?

Done that way. with_bridge no longer errors on a zero grant: the pool is created with anchor_held unset and jvm_held at zero, and try_grow calls take_missing_anchor before the reserve step, which issues a standalone one byte acquire while the flag is unset and never touches the real request. Once the byte lands the flag is set, so the retry is one extra JNI call per grow only while Spark keeps declining. Two grows retrying at once both get a byte; the second to lock hands its own back and books nothing. Drop returns the anchor only when it is held. A bridge error at construction still fails construction, since that is a JNI failure rather than memory pressure.

The short-grant branch in with_bridge is gone with it, and const _: () = assert!(ANCHOR_BYTES == 1) next to the constant pins the reasoning that a declined anchor is a zero grant with nothing to hand back. The comment above the construction acquire now says why it needs no catch_unwind (nothing reserved yet to roll back, and try_unwrap_or_throw catches the panic at the JNI boundary), and the retry carries the same note.

Tests: the declined-at-setup case now asserts the pool exists holding nothing, reports a short grant while the sibling still holds the share, takes the byte on the first grow after the share frees, makes exactly one JVM call per grow from then on, and returns the byte at drop; a second test pins that a pool that never held the anchor releases nothing on drop; a third drives two retries into the stub's acquire gate before either wins and checks one anchor is kept and one byte returned. Thirty loops in debug and release, no failures.

On #5466: left a note there proposing that whoever lands second rebases, with the per-reservation basis applied to the reserve-under-lock step rather than replacing it.

andygrove

The accounting audit comes out clean, and I went at this one hard because it is the area where a mistake is invisible until production. All three failure paths out of the bridge acquire roll used back through finish_acquire, jvm_held is only credited for bytes Spark actually returned so the short-grant release nets to zero, debit cannot underflow because used - jvm_held is exactly in_flight - anchor, and Drop returns the anchor exactly once. Over-commit is answered by construction since the limit check and state.used += additional are in the same critical section. The change also removes a latent re-entrancy deadlock on main where a spill callback could re-enter the pool while the mutex was held across the acquire.

Two things about where the anchor is taken. The first is that with_bridge takes it before the pool exists, so every task now pays a blocking JVM acquire at plan creation even when the plan will never allocate native memory. create_memory_pool runs unconditionally in createPlan and fair_unified is the default off-heap pool, so a plan that is only a scan and a filter and a projection registers no DataFusion MemoryConsumer, never calls try_grow, and on main makes no JNI memory call at all. Two things follow from taking the byte anyway. If the executor pool is full and the task is below poolSize / (2 * numActiveTasks), ExecutionMemoryPool.acquireMemory parks on that one byte and plan creation blocks on the task thread for memory the plan was not going to use. And since numActiveTasks is memoryForTask.keys.size, a task holding only the anchor still counts, so maxMemoryPerTask sits pinned at maxPoolSize / numActiveTasks for every other task on the executor, where before a task that dropped to zero left the set. Now that take_missing_anchor exists, would it make sense to drop the acquire in with_bridge and let the first try_grow take the anchor? The exposure between the anchor attempt and the request is the same window you already accept when Spark declines the byte, and try_new and create_memory_pool could go back to handing the pool back directly.

The second is that take_missing_anchor runs at the top of try_grow before the fair limit check, so while the anchor is missing a grow the pool is going to reject locally still makes a blocking JVM call first. try_grow_beyond_fair_limit_fails_without_calling_spark pins that an over-limit grow never reaches Spark, but it only holds because the anchor is already held in that test, so nothing covers the anchorless state. That matters beyond the extra round trip, because the one byte request can park when the executor pool has nothing free and the task is under its minimum share, so a grow that should have returned ResourcesExhausted immediately and let DataFusion spill blocks instead, and DataFusion retries try_grow around each spill so it repeats. Would it work to move take_missing_anchor below the reserve-under-lock step and above the bridge acquire, rolling back with finish_acquire(0, additional) if it errors?

Two process notes. CI still has not run on any commit carrying this design. The last real results are from eb410e517, which predates the anchor redesign, the non-fatal declined anchor and the registry restructuring. And even green says little here, since the only suites setting spark.memory.offHeap.enabled are CometRuntimeShutdownSuite, CometPluginsSuite, CometSparkSessionExtensionsSuite and the Iceberg suites. Separately, one line in the description reads more strongly than the test supports. TaskMemoryManager.acquireExecutionMemory is synchronized (this) while releaseExecutionMemory is not, so the JVM serializes a task's acquires regardless and second_acquire_does_not_queue_behind_a_parked_one is a statement about the Rust pool only. The doc comment on the test says "inside the pool" and is fine, it is just the description line. That also explains why the release p99 improvement is the real result in your table.

…it check

dwsmith1983 · 2026-09-15T10:57:36Z

Now that take_missing_anchor exists, would it make sense to drop the acquire in with_bridge and let the first try_grow take the anchor? ... try_new and create_memory_pool could go back to handing the pool back directly.

Done. Creating the pool makes no JVM call, new returns the pool directly and create_memory_pool no longer unwraps it, so a scan-filter-project plan never touches the task memory manager and never counts as an active task there. Pinned by a test that a pool which never grows makes no acquire and no release, and the registry race test now asserts neither create calls Spark and only the survivor's first grow takes the byte.

Would it work to move take_missing_anchor below the reserve-under-lock step and above the bridge acquire, rolling back with finish_acquire(0, additional) if it errors?

Done, and the anchor attempt sits under the same catch_unwind as the request so a panic inside its JNI frame rolls the reservation back too. over_limit_grow_without_an_anchor_never_reaches_spark pins that an over-limit grow with no anchor makes zero JVM calls, and the vanished-entry test now drives the parked anchor request from a first grow and asserts the reservation is back to zero when it fails.

one line in the description reads more strongly than the test supports.

Reworded: the description now says the second-acquire test pins the Rust side only, since TaskMemoryManager serializes a task's acquires on its own monitor, and the setup paragraph describes the lazy anchor. Pool tests are at 36 (two new, two reshaped), all of memory_pools looped ten times in release with no failure, full core suite and clippy clean.

dwsmith1983 force-pushed the perf/fair-pool-lock-free branch from 196d709 to 626bc39 Compare September 2, 2026 02:13

sunchao reviewed Sep 2, 2026

View reviewed changes

dwsmith1983 requested a review from sunchao September 2, 2026 10:37

dwsmith1983 and others added 3 commits September 2, 2026 21:58

Merge branch 'main' into perf/fair-pool-lock-free

d4ecb49

Merge branch 'main' into perf/fair-pool-lock-free

c11b0d5

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

eb410e5

sunchao reviewed Sep 3, 2026

View reviewed changes

dwsmith1983 added 3 commits September 4, 2026 10:43

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

e4edd86

dwsmith1983 requested a review from sunchao September 4, 2026 04:23

dwsmith1983 added 7 commits September 5, 2026 00:53

Merge branch 'main' into perf/fair-pool-lock-free

2b3db49

Merge branch 'main' into perf/fair-pool-lock-free

3ddfe3c

Merge branch 'main' into perf/fair-pool-lock-free

66c8cd1

Merge branch 'main' into perf/fair-pool-lock-free

67f7c3c

Merge branch 'main' into perf/fair-pool-lock-free

a9fc656

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

f96777b

Merge branch 'main' into perf/fair-pool-lock-free

1ed78a8

andygrove added enhancement New feature or request performance area:memory Memory pools, reservations, OOM handling labels Sep 6, 2026

dwsmith1983 added 5 commits September 7, 2026 20:14

Merge branch 'main' into perf/fair-pool-lock-free

2a9d0b3

Merge branch 'main' into perf/fair-pool-lock-free

e406b45

Merge branch 'main' into perf/fair-pool-lock-free

d1a3461

Merge branch 'main' into perf/fair-pool-lock-free

43f4468

Merge branch 'main' into perf/fair-pool-lock-free

270ca69

Merge branch 'main' into perf/fair-pool-lock-free

b7007ea

andygrove reviewed Sep 9, 2026

View reviewed changes

dwsmith1983 added 2 commits September 9, 2026 20:27

Merge branch 'main' into perf/fair-pool-lock-free

ea8f03b

Merge branch 'main' into perf/fair-pool-lock-free

14ef30e

andygrove reviewed Sep 9, 2026

View reviewed changes

dwsmith1983 added 2 commits September 10, 2026 08:07

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

d252b9f

dwsmith1983 mentioned this pull request Sep 10, 2026

fix: enforce fair memory share per reservation #5466

Open

dwsmith1983 requested a review from andygrove September 10, 2026 12:17

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

c281e73

sunchao mentioned this pull request Sep 10, 2026

fix: account fair-pool memory across sibling reservations #5847

Draft

dwsmith1983 added 2 commits September 11, 2026 07:17

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

1a94fba

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

b2aa4d5

andygrove mentioned this pull request Sep 11, 2026

Comet 1.1.0 Release (September) #5327

Open

dwsmith1983 added 3 commits September 12, 2026 00:10

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

3a6a8bb

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

c290167

Merge remote-tracking branch 'origin/main' into perf/fair-pool-lock-free

da77e26

andygrove reviewed Sep 13, 2026

View reviewed changes

dwsmith1983 added 3 commits September 14, 2026 21:51

Merge branch 'main' into perf/fair-pool-lock-free

9bd8869

Merge branch 'main' into perf/fair-pool-lock-free

724056b

perf: take the fair pool anchor on the first grow, after the fair lim…

e064d89

…it check

dwsmith1983 added 3 commits September 15, 2026 22:00

Merge branch 'main' into perf/fair-pool-lock-free

23a5220

Merge branch 'main' into perf/fair-pool-lock-free

7b371b1

Merge branch 'main' into perf/fair-pool-lock-free

cb87718

Uh oh!

Conversation

dwsmith1983 commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

sunchao Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 commented Sep 3, 2026

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

sunchao Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 commented Sep 4, 2026

Uh oh!

dwsmith1983 commented Sep 9, 2026

Uh oh!

andygrove left a comment

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 commented Sep 9, 2026

Uh oh!

dwsmith1983 commented Sep 9, 2026

Uh oh!

andygrove left a comment

Choose a reason for hiding this comment

A declined anchor now fails the task instead of spilling

Uh oh!

dwsmith1983 commented Sep 10, 2026

Uh oh!

andygrove left a comment

Choose a reason for hiding this comment

Uh oh!

dwsmith1983 commented Sep 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dwsmith1983 commented Sep 1, 2026 •

edited

Loading